{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "edd17e86-9d38-40a1-a9ea-5126b6d0073e",
   "metadata": {},
   "source": [
    "# Recall Complex Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9813037a-234d-4d56-aed4-a1f763c59a6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-11-04 19:37:26.266181: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "import openai\n",
    "import re\n",
    "import json\n",
    "import spacy\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b8faac47-5199-4e94-86d8-46c622f96d72",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = json.load(open('MathQA/train.json'))\n",
    "dev_data = json.load(open('MathQA/dev.json'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "41e59ee2-3ddc-4573-87fa-d343893c97ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dev_200_idx = np.random.choice(len(dev_data), 200, replace=False)\n",
    "# np.save('MathQA/dev_200_idx', dev_200_idx)\n",
    "\n",
    "dev_200_idx = np.load('MathQA/dev_200_idx.npy')\n",
    "dev_data_200 = []\n",
    "for i in dev_200_idx: dev_data_200.append(dev_data[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d8b7ff0-a089-4110-9c8a-0d7909761188",
   "metadata": {},
   "source": [
    "# Seed data as prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0a0bb752-2bd7-439a-9d41-a3d0b9a32734",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Problem': \"the banker ' s gain of a certain sum due 3 years hence at 10 % per annum is rs . 36 . what is the present worth ?\",\n",
       " 'Rationale': '\"explanation : t = 3 years r = 10 % td = ( bg × 100 ) / tr = ( 36 × 100 ) / ( 3 × 10 ) = 12 × 10 = rs . 120 td = ( pw × tr ) / 100 ⇒ 120 = ( pw × 3 × 10 ) / 100 ⇒ 1200 = pw × 3 pw = 1200 / 3 = rs . 400 answer : option a\"',\n",
       " 'options': 'a ) rs . 400 , b ) rs . 300 , c ) rs . 500 , d ) rs . 350 , e ) none of these',\n",
       " 'correct': 'a',\n",
       " 'annotated_formula': 'divide(multiply(const_100, divide(multiply(36, const_100), multiply(3, 10))), multiply(3, 10))',\n",
       " 'linear_formula': 'multiply(n2,const_100)|multiply(n0,n1)|divide(#0,#1)|multiply(#2,const_100)|divide(#3,#1)|',\n",
       " 'category': 'gain'}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "673b79cb-6254-4af2-93f6-802f71a27819",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Problem': 'a certain company ’ s profit in 1996 was 15 percent greater than its profit in 1995 , and its profit in 1997 was 35 percent greater than its profit in 1996 . the company ’ s profit in 1997 was what percent greater than its profit in 1995 ?',\n",
       " 'Rationale': '\"profit in 1995 - 100 profit in 1996 - 115 % increae profit in 1997 in comparison to 1995 = 15 + 115 * 35 % = 55.25 correct option : e\"',\n",
       " 'options': 'a ) 5 % , b ) 18 % , c ) 33 % , d ) 35 % , e ) 55.25 %',\n",
       " 'correct': 'e',\n",
       " 'annotated_formula': 'multiply(subtract(multiply(add(divide(35, const_100), const_1), add(const_1, divide(15, const_100))), const_1), const_100)',\n",
       " 'linear_formula': 'divide(n4,const_100)|divide(n1,const_100)|add(#0,const_1)|add(#1,const_1)|multiply(#2,#3)|subtract(#4,const_1)|multiply(#5,const_100)|',\n",
       " 'category': 'gain'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "716b5a8f-dec7-43ba-affc-00f88da91018",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Problem': 'how many positive integers less than 50 are multiples of 6 but not multiples of 8 ?',\n",
       " 'Rationale': '\"the lcm of 6 and 8 is 24 . if x < 50 and x is divisible by 6 not by 8 - - > x is not divisible by 24 . from 1 - - > 50 , we have 2 numbers which is divisible by 24 : 24 , 48 . from 1 - - > 50 , we have ( 48 - 6 ) / 6 + 1 = 8 numbers divisible by 6 . therefore , our answer is 8 - 2 = 6 numbers . b\"',\n",
       " 'options': 'a ) 4 , b ) 6 , c ) 8 , d ) 10 , e ) 12',\n",
       " 'correct': 'b',\n",
       " 'annotated_formula': 'divide(factorial(subtract(add(const_4, 6), const_1)), multiply(factorial(6), factorial(subtract(const_4, const_1))))',\n",
       " 'linear_formula': 'add(n1,const_4)|factorial(n1)|subtract(const_4,const_1)|factorial(#2)|subtract(#0,const_1)|factorial(#4)|multiply(#1,#3)|divide(#5,#6)|',\n",
       " 'category': 'general'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[900]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4fa536b9-4d0e-4a1d-90d8-6bd2d8e84da8",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Problem': 'the average of first 6 natural numbers is ?',\n",
       " 'Rationale': '\"sum of 6 natural no . = 42 / 2 = 21 average = 21 / 6 = 3.5 answer : b\"',\n",
       " 'options': 'a ) 5.2 , b ) 3.5 , c ) 5.3 , d ) 5.9 , e ) 5.1',\n",
       " 'correct': 'b',\n",
       " 'annotated_formula': 'add(6, const_1)',\n",
       " 'linear_formula': 'add(n0,const_1)|',\n",
       " 'category': 'general'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[1500]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ccf73cab-5241-4541-a072-d7db520b3344",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Problem': 'a and b together can do a work in 5 days . if a alone can do it in 15 days . in how many days can b alone do it ?',\n",
       " 'Rationale': '\"1 / 5 – 1 / 15 = 2 / 15 = > 7.5 answer : c\"',\n",
       " 'options': 'a ) 10 , b ) 99 , c ) 7.5 , d ) 55 , e ) 21',\n",
       " 'correct': 'c',\n",
       " 'annotated_formula': 'add(inverse(subtract(divide(const_1, 5), divide(const_1, 15))), divide(const_2, add(const_2, const_3)))',\n",
       " 'linear_formula': 'add(const_2,const_3)|divide(const_1,n0)|divide(const_1,n1)|divide(const_2,#0)|subtract(#1,#2)|inverse(#4)|add(#3,#5)|',\n",
       " 'category': 'physics'}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[3000]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa1aadff-8f73-43a6-9423-9948b95d4695",
   "metadata": {},
   "source": [
    "# Given a question, first paraphrase it to be 8 similar question, then use the seed complex CoT to generate the answer, then use them to prompt the target question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8eaee854-ec9d-4241-b670-2f3cd7ba6da9",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_paraphrase = open('MathQA/prompt_paraphrase.txt').read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "cb2a0503-eb1b-4003-a90e-516c9fa0abe8",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Given the question:  the banker ' s gain of a certain sum due 3 years hence at 10 % per annum is rs . 36 . what is the present worth ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "the banker ' s gain of a certain sum due 3 years hence at 10 % per annum is rs . 36 . During the next 3 years, he lost 5% per year. Later he made a fortune about rs. 10. what is the present worth ?\n",
      "[end of recall]\n",
      "\n",
      "Given the question:  a certain company ’ s profit in 1996 was 15 percent greater than its profit in 1995 , and its profit in 1997 was 35 percent greater than its profit in 1996 . the company ’ s profit in 1997 was what percent greater than its profit in 1995 ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "There is a company whose profit in 1993 was 10 percent greater than its profit in 1996, then its profit in 1996 was 15 percent greater than its profit in 1995 , and its profit in 1997 was 35 percent greater than its profit in 1993 . Then the company lost 10 million in 1998, which is of 10 percent of its profit in 1997. the company ’ s profit in 1997 was how much greater than its profit in 1995 ?\n",
      "[end of recall]\n",
      "\n",
      "Given the question:  how many positive integers less than 50 are multiples of 6 but not multiples of 8 ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "how many integers less than 500 and greater than 100 are multiples of 5 and 6 but not multiples of 8 ?\n",
      "[end of recall]\n",
      "\n",
      "Given the question: the average of first 6 natural numbers is ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "Compute the answer for the average of first 10 natural numbers plus the average of the first 10 primes\n",
      "[end of recall]\n"
     ]
    }
   ],
   "source": [
    "print(prompt_paraphrase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "cbd3260d-6d16-44c2-801a-67be4ef50271",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_q = prompt_paraphrase + \"\\n\\nGiven the question: \" + dev_data_200[0]['Problem'] + '\\n'\n",
    "prompt_q += \"Did you recall some similar yet longer question?\\n\"\n",
    "prompt_q += \"Answer: yes, I can recall a question as:\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "a65db6cf-f397-478f-9776-7d95f05caff6",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Given the question:  the banker ' s gain of a certain sum due 3 years hence at 10 % per annum is rs . 36 . what is the present worth ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "the banker ' s gain of a certain sum due 3 years hence at 10 % per annum is rs . 36 . During the next 3 years, he lost 5% per year. Later he made a fortune about rs. 10. what is the present worth ?\n",
      "[end of recall]\n",
      "\n",
      "Given the question:  a certain company ’ s profit in 1996 was 15 percent greater than its profit in 1995 , and its profit in 1997 was 35 percent greater than its profit in 1996 . the company ’ s profit in 1997 was what percent greater than its profit in 1995 ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "There is a company whose profit in 1993 was 10 percent greater than its profit in 1996, then its profit in 1996 was 15 percent greater than its profit in 1995 , and its profit in 1997 was 35 percent greater than its profit in 1993 . Then the company lost 10 million in 1998, which is of 10 percent of its profit in 1997. the company ’ s profit in 1997 was how much greater than its profit in 1995 ?\n",
      "[end of recall]\n",
      "\n",
      "Given the question:  how many positive integers less than 50 are multiples of 6 but not multiples of 8 ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "how many integers less than 500 and greater than 100 are multiples of 5 and 6 but not multiples of 8 ?\n",
      "[end of recall]\n",
      "\n",
      "Given the question: the average of first 6 natural numbers is ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "Compute the answer for the average of first 10 natural numbers plus the average of the first 10 primes\n",
      "[end of recall]\n",
      "\n",
      "Given the question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n"
     ]
    }
   ],
   "source": [
    "print(prompt_q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "f01340b1-19c7-4604-b149-e71be0338064",
   "metadata": {},
   "outputs": [],
   "source": [
    "q_neighbors_ = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                    prompt=prompt_q, \n",
    "                                    temperature=0.5, \n",
    "                                    max_tokens=256,\n",
    "                                    n=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "b9c842ab-1fc0-41dc-a6be-7f1cacba7261",
   "metadata": {},
   "outputs": [],
   "source": [
    "q_neighbors = []\n",
    "for q in q_neighbors_['choices']:\n",
    "    q = q['text'].split('[end of recall]')[0].strip()\n",
    "    q_neighbors.append(q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "f7ea359d-5cd4-4850-9e82-4ea1d2d680b1",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the journey from a to b is of 100 km.',\n",
       " 'a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey?',\n",
       " 'a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey?',\n",
       " 'A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey ?',\n",
       " 'a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . on the way back, he takes a break for 15 minutes . find his average speed for the entire journey ?',\n",
       " 'a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the total journey is 150 km ?',\n",
       " 'a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the total journey is of 10 hours.',\n",
       " 'A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the journey from a to b is of 120 km ?']"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q_neighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "fdc9c46b-39f2-45fa-a156-33eadfae5c4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_complex = open('MathQA/prompt_complex_no_option.txt').read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "04972d26-5d67-4b3f-9046-7620e629498b",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: there were 35 students in a hostel . due to the admission of 7 new students the expenses of the mess were increased by rs . 84 per day while the average expenditure per head diminished by re 1 . what was the original expenditure of the mess ?\n",
      "Let's think step by step\n",
      "let the original average expenditure be rs . x\n",
      "then , 42 ( x - 1 ) - 35 x = 84\\,7 x = 126\n",
      "x = 18\n",
      "therefore original expenditure = rs . ( 35 * 18 ) = rs . 630.\n",
      "The answer is rs . 630\n",
      "\n",
      "Question: a train 200 m long passes a man , running at 5 km / hr in the same direction in which the train is going , in 10 seconds . the speed of the train is ?\n",
      "Let's think step by step\n",
      "speed of the train relative to man = ( 200 / 10 ) m / sec = ( 20 ) m / sec.\n",
      "[ ( 20 ) * ( 18 / 5 ) ] km / hr = 72 km / hr.\n",
      "let the speed of the train be x km / hr.\n",
      "then , relative speed = ( x - 5 ) km / hr.\n",
      "x - 5 = 72, x = 77 km / hr .\n",
      "The answer is 77\n",
      "\n",
      "Question: solution x contains 20 % of material a and 80 % of material b . solution y contains 30 % of material a and 70 % of material b . a mixture of both these solutions contains 22 % of material a in the final product . how much solution x is present in the mixture ?\n",
      "Let's think step by step\n",
      "we can assume the total weight of the mixture = 100\n",
      "conc of a in the final mixture = 22\n",
      "let weight of a in the mixture be x.\n",
      "conc given = 20% = 0.2\n",
      "therefore , weight of b = 100 - x.\n",
      "conc given = 30% = 0.3\n",
      "now , accordding to the problem, 0.2 x + 0.3 ( 100 - x ) = 22\n",
      "solving , we get x = 80\n",
      "since we assumed the weight of the mixture = 100, therefore presence of a in the mixture = 80%.\n",
      "The answer is 80%\n",
      "\n",
      "Question: a trader sells 40 metres of cloth for rs . 8200 at a profit of rs . 35 per metre of cloth . how much profit will the trder earn on 40 metres of cloth ?\n",
      "Let's think step by step\n",
      "price of 1 metre cloth = 8200 / 40 = rs 205\n",
      "cost of 1 metre cloth = rs 205 – 35 = rs 170\n",
      "cost on 40 metres = 170 x 40 = rs . 6800\n",
      "profit earned on 40 metres cloth = rs . 8200 – rs . 6800 = rs . 1400\n",
      "The answer is rs . 1400\n",
      "\n",
      "Question: if x < y < z and y - x > 5 , where x is an even integer and y and z are odd integers , what is the least possible value s of z - x ?\n",
      "Let's think step by step\n",
      "We know x < y < z\n",
      "to find the least possible value for z - x, we need to find the values for z and x that can be closest to each other.\n",
      "if x is some even number, then what could be minimum possible odd z.\n",
      "if x is some even number, y - x > 5 ; y > x + 5\n",
      "minimum value for y = x + 5 + 2 = x + 7\n",
      "(note : x + 5 is as even + odd = odd and nearest odd greater than x + 5 is x + 5 + 2)\n",
      "minimum value for z = y + 2 = x + 7 + 2 = x + 9\n",
      "(note : z = y + 2 because both z and y are odd, difference between two odd numbers is 2)\n",
      "s = z - x = x + 9 - x = 9\n",
      "The answer is 9\n",
      "\n",
      "Question: what is the difference between the c . i . on rs . 6000 for 1 1 / 2 years at 4 % per annum compounded yearly and half - yearly ?\n",
      "Let's think step by step\n",
      "c . i . when interest is compounded yearly = [ 6000 * ( 1 + 4 / 100 ) * ( 1 + ( 1 / 2 * 4 ) / 100 ] = 6000 * 26 / 25 * 51 / 50 = rs . 6364.8\n",
      "c . i . when interest is compounded half - yearly = [ 6000 * ( 1 + 2 / 100 ) 2 ] = ( 6000 * 51 / 50 * 51 / 50 * 51 / 50 ) = rs . 6367.25\n",
      "difference = ( 6367.25 - 6364.8 ) = rs . 2.45.\n",
      "The answer is s . 2.45\n",
      "\n",
      "Question: the average weight of a , b and c is 45 kg . if the average weight of a and b be 40 kg and that of b and c be 45 kg , then the weight of b is :\n",
      "Let's think step by step\n",
      "let a , b , c represent their respective weights.\n",
      "then , we have : a + b + c = ( 45 x 3 ) = 135 … ( i )\n",
      "a + b = ( 40 x 2 ) = 80 … ( ii )\n",
      "b + c = ( 45 x 2 ) = 90 … ( iii )\n",
      "adding ( ii ) and ( iii ), we get : a + 2 b + c = 170 … ( iv )\n",
      "subtracting ( i ) from ( iv ), we get : b = 35.\n",
      "b’ s weight = 35 kg .\n",
      "The answer is 35 kg\n",
      "\n",
      "Question: the compound and the simple interests on a certain sum at the same rate of interest for two years are rs . 11730 and rs . 10200 respectively . find the sum\n",
      "Let's think step by step\n",
      "the simple interest for the first year is 10200 / 2 is rs . 5100\n",
      "and compound interest for first year also is rs . 5100.\n",
      "the compound interest for second year on rs . 5100 for one year\n",
      "so rate of the interest = ( 100 * 1530 ) / ( 5100 * 1 ) = 30% p\n",
      "so p = ( 100 * 10200 ) / ( 30 * 2 ) = rs . 17000\n",
      "The answer is rs . 17000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prompt_complex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "3ab6f715-6a38-4861-ba0e-76b4ad5fa605",
   "metadata": {},
   "outputs": [],
   "source": [
    "rationales = []\n",
    "for q in q_neighbors:\n",
    "    prompt_q = '' \n",
    "    prompt_q += 'Question: ' + q + '\\n'\n",
    "    prompt_q += \"Let's think step by step\\n\"\n",
    "    prompt_q = prompt_complex + '\\n\\n' + prompt_q\n",
    "    response = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                        prompt=prompt_q, \n",
    "                                        temperature=0.5, \n",
    "                                        max_tokens=256)\n",
    "    rationales.append(response['choices'][0]['text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "205fb23b-154f-4e69-9e3d-75b11ce91056",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Let's take the journey from a to b as 1 hour.\n",
      "\n",
      "Therefore, the man covers 20 km in 1 hour.\n",
      "\n",
      "In 2 hours, he covers 20 + 30 = 50 km.\n",
      "\n",
      "Therefore, his average speed for the journey is (20 + 50)/2 = 35 kmph.\n",
      "\n",
      "----\n",
      "the average speed for the first half of the journey = ( 20 + 30 ) / 2 = 25 kmph\n",
      "the average speed for the second half of the journey = ( 20 + 30 ) / 2 = 25 kmph\n",
      "The answer is 25 kmph\n",
      "\n",
      "----\n",
      "distance = speed * time\n",
      "distance from a to b = 20 * t\n",
      "distance from b to a = 30 * t\n",
      "since the total distance is same, therefore 20 * t = 30 * t\n",
      "t = 0\n",
      "average speed for the first half of the journey = 20 kmph\n",
      "average speed for the second half of the journey = 30 kmph\n",
      "\n",
      "----\n",
      "let the total distance be d km.\n",
      "time taken to go from a to b = d / 20 hr\n",
      "time taken to come back from b to a = d / 30 hr\n",
      "total time taken for the journey = d / 20 + d / 30 hr\n",
      "average speed for the first half of the journey = ( d / 2 ) / ( d / 20 ) = 10 kmph\n",
      "average speed for the second half of the journey = ( d / 2 ) / ( d / 30 ) = 15 kmph.\n",
      "The answer is 10 kmph and 15 kmph\n",
      "\n",
      "----\n",
      "the man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . on the way back, he takes a break for 15 minutes . find his average speed for the entire journey ?\n",
      "let's take the entire journey to be of 1 hour.\n",
      "then distance covered in 1 hour = ( 20 + 30 ) km = 50 km\n",
      "time taken for the entire journey = ( 1 + ( 1 / 4 ) ) hour = 1.25 hour\n",
      "average speed = ( 50 / 1.25 ) kmph = 40 kmph\n",
      "The answer is 40 kmph\n",
      "\n",
      "----\n",
      "let the average speed be x kmph.\n",
      "then , ( 20 + 30 ) / 2 = x\n",
      "x = 25 kmph.\n",
      "The answer is 25 kmph\n",
      "\n",
      "----\n",
      "distance from a to b = 20 * 10 = 200 km\n",
      "distance from b to a = 30 * 10 = 300 km\n",
      "total distance = 200 + 300 = 500 km\n",
      "average speed = 500 / 10 = 50 kmph\n",
      "The answer is 50 kmph\n",
      "\n",
      "----\n",
      "let the average speed of the entire journey be x kmph.\n",
      "then , ( x - 20 ) * 2 + ( x - 30 ) * 2 = 120\n",
      "solving , we get x = 30 kmph.\n",
      "The answer is 30\n",
      "\n",
      "----\n"
     ]
    }
   ],
   "source": [
    "for r in rationales: \n",
    "    print(r + '\\n\\n----')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "84bebc44-3b4e-4c8a-8b82-3b08a10bf91b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def construct_options(rationale):\n",
    "    pattern = '\\d*\\.?\\d+'\n",
    "    ans_ = re.findall(pattern, rationale)\n",
    "    if(len(ans_) >= 1): \n",
    "        ans_ = ans_[-1]\n",
    "        unit = rationale.split(ans_)[-1].strip().replace('.', '')\n",
    "        if('.' in ans_): ans_ = float(ans_)\n",
    "        else: ans_ = int(ans_)\n",
    "    else: \n",
    "        ans_ = 18\n",
    "        unit = ''\n",
    "        \n",
    "    c = np.random.uniform()\n",
    "    rint = np.random.choice([-1, -2, -3, -4, -5, 1, 2, 3, 4, 5], 3, replace=False)\n",
    "    choices = ''\n",
    "    k = 0\n",
    "    for i, ci in enumerate(['a', 'b', 'c', 'd']):\n",
    "        if(i <= 4 * c and 4 * c <= i + 1):\n",
    "            choices += ci + ' ) ' + str(ans_)\n",
    "            if(unit != ''): choices += ' ' + unit\n",
    "            a = ci + ', ' + str(ans_)\n",
    "            if(unit != ''): a += ' ' + unit\n",
    "        else: \n",
    "            choices += ci + ' ) ' + str(ans_ + rint[k])\n",
    "            if(unit != ''): choices += ' ' + unit\n",
    "            k += 1\n",
    "        if(i < 3): choices += ' , '\n",
    "        \n",
    "    # print(choices)\n",
    "    rationale_ = rationale.split('\\n')\n",
    "    if('the answer is' in rationale_[-1].lower()):\n",
    "        rationale = list(rationale_[:-1])\n",
    "        ans_str = 'The answer is ' + a\n",
    "        rationale.append(ans_str)\n",
    "        rationale = '\\n'.join(rationale)\n",
    "    else: \n",
    "        rationale = list(rationale_)\n",
    "        ans_str = 'The answer is ' + a\n",
    "        rationale.append(ans_str)\n",
    "        rationale = '\\n'.join(rationale)\n",
    "    return rationale, choices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "af09d3eb-beb0-4ef7-b6bc-bd677a32ed42",
   "metadata": {},
   "outputs": [],
   "source": [
    "rationale_, choices = construct_options(rationales[3].replace('\\n\\n', '\\n'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "c6b4f038-31d7-4372-971e-31fd0d828fe8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a ) 17 kmph , b ) 12 kmph , c ) 15 kmph , d ) 11 kmph\n",
      " let the total distance be d km.\n",
      "time taken to go from a to b = d / 20 hr\n",
      "time taken to come back from b to a = d / 30 hr\n",
      "total time taken for the journey = d / 20 + d / 30 hr\n",
      "average speed for the first half of the journey = ( d / 2 ) / ( d / 20 ) = 10 kmph\n",
      "average speed for the second half of the journey = ( d / 2 ) / ( d / 30 ) = 15 kmph.\n",
      "The answer is c, 15 kmph\n"
     ]
    }
   ],
   "source": [
    "print(choices + '\\n', rationale_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "2c01e492-1b8c-4fe9-8550-7f2b60475135",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_final = ''\n",
    "for q, r in zip(q_neighbors, rationales):\n",
    "    r_, choices = construct_options(r.strip().replace('\\n\\n', '\\n'))\n",
    "    prompt_final += 'Question: ' + q + '\\n'\n",
    "    prompt_final += 'Options: ' + choices + '\\n'\n",
    "    prompt_final += \"Let's think step by step\\n\"\n",
    "    prompt_final += r_ + '\\n\\n'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "1d1bca19-2c2d-4c63-a9ac-a46eae4a0d5c",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the journey from a to b is of 100 km.\n",
      "Options: a ) 35 kmph , b ) 37 kmph , c ) 36 kmph , d ) 32 kmph\n",
      "Let's think step by step\n",
      "Let's take the journey from a to b as 1 hour.\n",
      "Therefore, the man covers 20 km in 1 hour.\n",
      "In 2 hours, he covers 20 + 30 = 50 km.\n",
      "Therefore, his average speed for the journey is (20 + 50)/2 = 35 kmph.\n",
      "The answer is a, 35 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey?\n",
      "Options: a ) 27 kmph , b ) 23 kmph , c ) 20 kmph , d ) 25 kmph\n",
      "Let's think step by step\n",
      "the average speed for the first half of the journey = ( 20 + 30 ) / 2 = 25 kmph\n",
      "the average speed for the second half of the journey = ( 20 + 30 ) / 2 = 25 kmph\n",
      "The answer is d, 25 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey?\n",
      "Options: a ) 25 kmph , b ) 26 kmph , c ) 29 kmph , d ) 30 kmph\n",
      "Let's think step by step\n",
      "distance = speed * time\n",
      "distance from a to b = 20 * t\n",
      "distance from b to a = 30 * t\n",
      "since the total distance is same, therefore 20 * t = 30 * t\n",
      "t = 0\n",
      "average speed for the first half of the journey = 20 kmph\n",
      "average speed for the second half of the journey = 30 kmph\n",
      "The answer is d, 30 kmph\n",
      "\n",
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey ?\n",
      "Options: a ) 11 kmph , b ) 15 kmph , c ) 19 kmph , d ) 12 kmph\n",
      "Let's think step by step\n",
      "let the total distance be d km.\n",
      "time taken to go from a to b = d / 20 hr\n",
      "time taken to come back from b to a = d / 30 hr\n",
      "total time taken for the journey = d / 20 + d / 30 hr\n",
      "average speed for the first half of the journey = ( d / 2 ) / ( d / 20 ) = 10 kmph\n",
      "average speed for the second half of the journey = ( d / 2 ) / ( d / 30 ) = 15 kmph.\n",
      "The answer is b, 15 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . on the way back, he takes a break for 15 minutes . find his average speed for the entire journey ?\n",
      "Options: a ) 40 kmph , b ) 41 kmph , c ) 45 kmph , d ) 39 kmph\n",
      "Let's think step by step\n",
      "the man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . on the way back, he takes a break for 15 minutes . find his average speed for the entire journey ?\n",
      "let's take the entire journey to be of 1 hour.\n",
      "then distance covered in 1 hour = ( 20 + 30 ) km = 50 km\n",
      "time taken for the entire journey = ( 1 + ( 1 / 4 ) ) hour = 1.25 hour\n",
      "average speed = ( 50 / 1.25 ) kmph = 40 kmph\n",
      "The answer is a, 40 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the total journey is 150 km ?\n",
      "Options: a ) 25 kmph , b ) 27 kmph , c ) 30 kmph , d ) 21 kmph\n",
      "Let's think step by step\n",
      "let the average speed be x kmph.\n",
      "then , ( 20 + 30 ) / 2 = x\n",
      "x = 25 kmph.\n",
      "The answer is a, 25 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the total journey is of 10 hours.\n",
      "Options: a ) 48 kmph , b ) 49 kmph , c ) 50 kmph , d ) 46 kmph\n",
      "Let's think step by step\n",
      "distance from a to b = 20 * 10 = 200 km\n",
      "distance from b to a = 30 * 10 = 300 km\n",
      "total distance = 200 + 300 = 500 km\n",
      "average speed = 500 / 10 = 50 kmph\n",
      "The answer is c, 50 kmph\n",
      "\n",
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the journey from a to b is of 120 km ?\n",
      "Options: a ) 27 , b ) 30 , c ) 31 , d ) 25\n",
      "Let's think step by step\n",
      "let the average speed of the entire journey be x kmph.\n",
      "then , ( x - 20 ) * 2 + ( x - 30 ) * 2 = 120\n",
      "solving , we get x = 30 kmph.\n",
      "The answer is b, 30\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prompt_final)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "4882decd-a44e-4456-8bb3-c560e5a18b2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_final += 'Question: ' + dev_data_200[0]['Problem'] + '\\n'\n",
    "prompt_final += 'Options: ' + dev_data_200[0]['options'] + '\\n'\n",
    "prompt_final += \"Let's think step by step\\n\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "c7488deb-fd16-49f4-86fd-753777eb7f58",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the journey from a to b is of 100 km.\n",
      "Options: a ) 35 kmph , b ) 37 kmph , c ) 36 kmph , d ) 32 kmph\n",
      "Let's think step by step\n",
      "Let's take the journey from a to b as 1 hour.\n",
      "Therefore, the man covers 20 km in 1 hour.\n",
      "In 2 hours, he covers 20 + 30 = 50 km.\n",
      "Therefore, his average speed for the journey is (20 + 50)/2 = 35 kmph.\n",
      "The answer is a, 35 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey?\n",
      "Options: a ) 27 kmph , b ) 23 kmph , c ) 20 kmph , d ) 25 kmph\n",
      "Let's think step by step\n",
      "the average speed for the first half of the journey = ( 20 + 30 ) / 2 = 25 kmph\n",
      "the average speed for the second half of the journey = ( 20 + 30 ) / 2 = 25 kmph\n",
      "The answer is d, 25 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey?\n",
      "Options: a ) 25 kmph , b ) 26 kmph , c ) 29 kmph , d ) 30 kmph\n",
      "Let's think step by step\n",
      "distance = speed * time\n",
      "distance from a to b = 20 * t\n",
      "distance from b to a = 30 * t\n",
      "since the total distance is same, therefore 20 * t = 30 * t\n",
      "t = 0\n",
      "average speed for the first half of the journey = 20 kmph\n",
      "average speed for the second half of the journey = 30 kmph\n",
      "The answer is d, 30 kmph\n",
      "\n",
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey ?\n",
      "Options: a ) 11 kmph , b ) 15 kmph , c ) 19 kmph , d ) 12 kmph\n",
      "Let's think step by step\n",
      "let the total distance be d km.\n",
      "time taken to go from a to b = d / 20 hr\n",
      "time taken to come back from b to a = d / 30 hr\n",
      "total time taken for the journey = d / 20 + d / 30 hr\n",
      "average speed for the first half of the journey = ( d / 2 ) / ( d / 20 ) = 10 kmph\n",
      "average speed for the second half of the journey = ( d / 2 ) / ( d / 30 ) = 15 kmph.\n",
      "The answer is b, 15 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . on the way back, he takes a break for 15 minutes . find his average speed for the entire journey ?\n",
      "Options: a ) 40 kmph , b ) 41 kmph , c ) 45 kmph , d ) 39 kmph\n",
      "Let's think step by step\n",
      "the man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . on the way back, he takes a break for 15 minutes . find his average speed for the entire journey ?\n",
      "let's take the entire journey to be of 1 hour.\n",
      "then distance covered in 1 hour = ( 20 + 30 ) km = 50 km\n",
      "time taken for the entire journey = ( 1 + ( 1 / 4 ) ) hour = 1.25 hour\n",
      "average speed = ( 50 / 1.25 ) kmph = 40 kmph\n",
      "The answer is a, 40 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the total journey is 150 km ?\n",
      "Options: a ) 25 kmph , b ) 27 kmph , c ) 30 kmph , d ) 21 kmph\n",
      "Let's think step by step\n",
      "let the average speed be x kmph.\n",
      "then , ( 20 + 30 ) / 2 = x\n",
      "x = 25 kmph.\n",
      "The answer is a, 25 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the total journey is of 10 hours.\n",
      "Options: a ) 48 kmph , b ) 49 kmph , c ) 50 kmph , d ) 46 kmph\n",
      "Let's think step by step\n",
      "distance from a to b = 20 * 10 = 200 km\n",
      "distance from b to a = 30 * 10 = 300 km\n",
      "total distance = 200 + 300 = 500 km\n",
      "average speed = 500 / 10 = 50 kmph\n",
      "The answer is c, 50 kmph\n",
      "\n",
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the journey from a to b is of 120 km ?\n",
      "Options: a ) 27 , b ) 30 , c ) 31 , d ) 25\n",
      "Let's think step by step\n",
      "let the average speed of the entire journey be x kmph.\n",
      "then , ( x - 20 ) * 2 + ( x - 30 ) * 2 = 120\n",
      "solving , we get x = 30 kmph.\n",
      "The answer is b, 30\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey ?\n",
      "Options: a ) 11 , b ) 24 , c ) 55 , d ) 66 , e ) 71\n",
      "Let's think step by step\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prompt_final)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "d1808e84-6fd5-4601-a60f-d8d0fb1e8114",
   "metadata": {},
   "outputs": [],
   "source": [
    "ans_final = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                    prompt=prompt_final, \n",
    "                                    temperature=0, \n",
    "                                    max_tokens=256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "ca7b8bc0-d93b-4190-8b02-c0edbef81c03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "distance from a to b = 20 * t\n",
      "distance from b to a = 30 * t\n",
      "total distance = 20 * t + 30 * t\n",
      "average speed = ( 20 * t + 30 * t ) / ( 2 * t )\n",
      "average speed = ( 50 * t ) / ( 2 * t )\n",
      "average speed = 25 kmph.\n",
      "The answer is 25 kmph\n"
     ]
    }
   ],
   "source": [
    "print(ans_final['choices'][0]['text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "269e284e-2e19-4d0f-a2c4-ded5571b39b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Problem': 'a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey ?',\n",
       " 'Rationale': '\"distance from a and b be \\' d \\' average speed = total distance / total time average speed = ( 2 d ) / [ ( d / 20 ) + ( d / 30 ) ] = ( 2 d ) / [ 5 d / 60 ) = > 24 kmph . answer : b\"',\n",
       " 'options': 'a ) 11 , b ) 24 , c ) 55 , d ) 66 , e ) 71',\n",
       " 'correct': 'b',\n",
       " 'annotated_formula': 'divide(add(20, 30), const_2)',\n",
       " 'linear_formula': 'add(n0,n1)|divide(#0,const_2)|',\n",
       " 'category': 'physics'}"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dev_data_200[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "id": "eac5a777-aa59-4fb7-bc94-4f6a766e9630",
   "metadata": {},
   "outputs": [],
   "source": [
    "def construct_prompt_recall(prompt_paraphrase, prompt_seed, q):\n",
    "    # construct paraphrase\n",
    "    prompt_q = prompt_paraphrase + \"\\n\\nGiven the question: \" + q + '\\n'\n",
    "    prompt_q += \"Did you recall some similar yet longer question?\\n\"\n",
    "    prompt_q += \"Answer: yes, I can recall a question as:\"\n",
    "    \n",
    "    q_neighbors_ = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                    prompt=prompt_q, \n",
    "                                    temperature=0.5, \n",
    "                                    max_tokens=256,\n",
    "                                    n=8)\n",
    "    \n",
    "    q_neighbors = []\n",
    "    for q in q_neighbors_['choices']:\n",
    "        q = q['text'].split('[end of recall]')[0].strip()\n",
    "        q_neighbors.append(q)\n",
    "    \n",
    "    # construct pseudo cot\n",
    "    rationales = []\n",
    "    for q in q_neighbors:\n",
    "        prompt_q = '' \n",
    "        prompt_q += 'Question: ' + q + '\\n'\n",
    "        prompt_q += \"Let's think step by step\\n\"\n",
    "        prompt_q = prompt_seed + '\\n\\n' + prompt_q\n",
    "        response = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                            prompt=prompt_q, \n",
    "                                            temperature=0.5, \n",
    "                                            max_tokens=256)\n",
    "        rationales.append(response['choices'][0]['text'])\n",
    "        \n",
    "    # compose neighbor question and rationale to prompt\n",
    "    prompt_final = ''\n",
    "    for q, r in zip(q_neighbors, rationales):\n",
    "        r_, choices = construct_options(r.strip().replace('\\n\\n', '\\n'))\n",
    "        prompt_final += 'Question: ' + q + '\\n'\n",
    "        prompt_final += 'Options: ' + choices + '\\n'\n",
    "        prompt_final += \"Let's think step by step\\n\"\n",
    "        prompt_final += r_ + '\\n\\n'\n",
    "    return prompt_final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "95cb3897-9f75-423a-bade-ddac10d4464c",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_recall = construct_prompt_recall(prompt_paraphrase, prompt_complex, dev_data_200[0]['Problem'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "4ad7db0e-ffed-49bb-b451-afb7ffac23b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the journey from a to b took 2 hours and the journey from b to a took 1.5 hours.\n",
      "Options: a ) 21.86 kmph , b ) 20.86 kmph , c ) 19.86 kmph , d ) 22.86 kmph\n",
      "Let's think step by step\n",
      "let the distance from a to b be d km.\n",
      "then , d / 20 = 2\\,d = 40 km\n",
      "speed for entire journey = ( 40 + 40 ) / ( 2 + 1.5 ) = 80 / 3.5 = 22.86 kmph.\n",
      "The answer is d, 22.86 kmph\n",
      "\n",
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . On the way, he stops at c for 10 minutes. Find his average speed for the entire journey ?\n",
      "Options: a ) 1800 kmph , b ) 1798 kmph , c ) 1804 kmph , d ) 1799 kmph\n",
      "Let's think step by step\n",
      "time taken to go from a to b = ( 20 * 60 ) / 60 = 20 minutes\n",
      "time taken to go from b to a = ( 30 * 60 ) / 60 = 30 minutes\n",
      "time taken to stop at c = 10 minutes\n",
      "total time taken for the journey = 20 + 30 + 10 = 60 minutes\n",
      "distance travelled = ( 20 * 60 ) + ( 30 * 60 ) = 1800 km\n",
      "average speed = ( 1800 * 60 ) / 60 = 1800 kmph\n",
      "The answer is a, 1800 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the total journey time is 1 hour?\n",
      "Options: a ) 2 kmph , b ) 1 kmph , c ) 4 kmph , d ) 6 kmph\n",
      "Let's think step by step\n",
      "let distance between a and b be d km.\n",
      "time taken to go from a to b = ( d / 20 ) hours\n",
      "time taken to come back from b to a = ( d / 30 ) hours\n",
      "total time taken for the journey = ( d / 20 ) + ( d / 30 ) = ( 6 d + 9 d ) / 60 = 15 d / 60\n",
      "d / ( 15 d / 60 ) = 60 / 15 = 4 kmph\n",
      "The answer is c, 4 kmph\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the first half of the journey and the second half of the journey?\n",
      "Options: a ) 13 , b ) 15 , c ) 20 , d ) 17\n",
      "Let's think step by step\n",
      "let the man's average speed for the first half of the journey be x kmph.\n",
      "then , his average speed for the second half of the journey = ( 2 x - 20 ) kmph.\n",
      "now , according to the question , ( x + ( 2 x - 20 ) ) / 2 = 25\n",
      "solving , we get x = 15 kmph.\n",
      "therefore , average speed for the first half of the journey = 15 kmph.\n",
      "The answer is b, 15\n",
      "\n",
      "Question: a man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the total journey is of 200 kms ?\n",
      "Options: a ) 98 kmph , b ) 97 kmph , c ) 100 kmph , d ) 96 kmph\n",
      "Let's think step by step\n",
      "let the speed of the man be x kmph.\n",
      "then , ( x - 20 ) kmph = 30\\,x = 50 kmph.\n",
      "distance covered in forward journey = ( 200 * 20 ) / 50 = 80 km.\n",
      "distance covered in return journey = ( 200 * 30 ) / 50 = 120 km.\n",
      "average speed for the entire journey = ( 80 + 120 ) / 2 = 100 kmph.\n",
      "The answer is c, 100 kmph\n",
      "\n",
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey ?\n",
      "\n",
      "Given the question: how many days are there in a year on planet x ?\n",
      "Did you recall some similar yet longer question?\n",
      "Answer: yes, I can recall a question as:\n",
      "How many days are there in a year on planet x if a day on planet x is 24 hours long?\n",
      "Options: a ) 26 hours in a year , b ) 28 hours in a year , c ) 24 hours in a year , d ) 21 hours in a year\n",
      "Let's think step by step\n",
      "There are 24 hours in a day.\n",
      "There are 365 days in a year.\n",
      "Therefore, there are 365 * 24 hours in a year.\n",
      "The answer is c, 24 hours in a year\n",
      "\n",
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph . find his average speed for the entire journey if the journey from a to b is of 120 km ?\n",
      "Options: a ) 120 kmph , b ) 124 kmph , c ) 123 kmph , d ) 125 kmph\n",
      "Let's think step by step\n",
      "average speed for entire journey = ( 120 + 120 ) / 2 = 120 kmph\n",
      "speed of the man from a to b = 20 kmph\n",
      "speed of the man from b to a = 30 kmph\n",
      "average speed for the journey from a to b = ( 20 + 30 ) / 2 = 25 kmph\n",
      "average speed for the journey from b to a = ( 30 + 20 ) / 2 = 25 kmph\n",
      "The answer is a, 120 kmph\n",
      "\n",
      "Question: A man goes from a to b at a speed of 20 kmph and comes back to a at a speed of 30 kmph. In the meantime, he finds a c and goes to c at a speed of 40 kmph and comes back to b at a speed of 60 kmph. Finally, he goes from b to a at a speed of 80 kmph. Find his average speed for the entire journey.\n",
      "Options: a ) 45 kmph , b ) 48 kmph , c ) 53 kmph , d ) 50 kmph\n",
      "Let's think step by step\n",
      "The man's average speed for the entire journey is:\n",
      "(20 + 30 + 40 + 60 + 80)/5 = 50 kmph\n",
      "The answer is d, 50 kmph\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prompt_recall)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a82b5d0d-7a7b-4ef8-8090-b43b2a4c4459",
   "metadata": {},
   "source": [
    "# Run -- Acc 23.5 v.s. complex acc 42.5, retrieval acc 69.5 -- So the performance is indeed correlated with annotation quality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "bbd07463-10b1-4844-bae3-ca2617c6ff0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_answer(pred_str, ans_str):\n",
    "    \"\"\"Find the last number as the predicted answer\"\"\"\n",
    "    pattern = '\\d*\\.?\\d+'\n",
    "    pred = re.findall(pattern, pred_str)\n",
    "    if(len(pred) >= 1):\n",
    "        # print(pred_str)\n",
    "        pred = pred[-1]\n",
    "        gold = re.findall(pattern, ans_str)\n",
    "        # print(ans_str)\n",
    "        gold = gold[-1]\n",
    "        return pred == gold\n",
    "    else: return False\n",
    "\n",
    "def parse_pred_ans(filename):\n",
    "    with open(filename) as fd: lines = fd.readlines()\n",
    "    am, a = None, None\n",
    "    num_q, acc = 0, 0\n",
    "    current_mode = 'none'\n",
    "    questions = []\n",
    "    ans_pred = []\n",
    "    ans_gold = []\n",
    "    for l in lines:\n",
    "        if(l.startswith('Question: ')):\n",
    "            if(am is not None and a is not None):\n",
    "                questions.append(q)\n",
    "                ans_pred.append(am)\n",
    "                ans_gold.append(a)\n",
    "                if(test_answer(am, a)):\n",
    "                    acc += 1\n",
    "            current_mode = 'q'\n",
    "            q = l\n",
    "            num_q += 1\n",
    "        elif(l.startswith('Options: ')): pass\n",
    "        elif(l.startswith('A_model:')):\n",
    "            current_mode = 'am'\n",
    "            am = l\n",
    "        elif(l.startswith('A:')):\n",
    "            current_mode = 'a'\n",
    "            a = l\n",
    "        else:\n",
    "            if(current_mode == 'q'): q += l\n",
    "            elif(current_mode == 'am'): am += l\n",
    "            elif(current_mode == 'a'): a += l\n",
    "            else:\n",
    "                raise ValueError(current_mode)\n",
    "                \n",
    "    questions.append(q)\n",
    "    ans_pred.append(am)\n",
    "    ans_gold.append(a)\n",
    "    if(test_answer(am, a)):\n",
    "        acc += 1\n",
    "    print('num_q %d correct %d ratio %.4f' % (num_q, acc, float(acc / num_q)))\n",
    "    return questions, ans_pred, ans_gold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "09d420a0-61ad-4122-9f69-16ae1506f3e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [2:24:42<00:00, 43.41s/it]\n"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "with open('outputs/dev_outputs_recall.txt', 'w') as fd, open('outputs/dev_recall_prompts.txt', 'w') as fd_prompts:\n",
    "    for case in tqdm(dev_data_200):\n",
    "        q = case['Problem']\n",
    "        c = case['options']\n",
    "        a = case['Rationale']\n",
    "        cr = case['correct']\n",
    "    \n",
    "        ans_dict = {}\n",
    "        for ac in c.split(', '):\n",
    "            idx_, num_ = ac.split(' ) ')[-2:]\n",
    "            if('a' in idx_): ans_dict['a'] = num_\n",
    "            if('b' in idx_): ans_dict['b'] = num_\n",
    "            if('c' in idx_): ans_dict['c'] = num_\n",
    "            if('d' in idx_): ans_dict['d'] = num_\n",
    "            if('e' in idx_): ans_dict['e'] = num_\n",
    "            \n",
    "        a += ' ' + ans_dict[cr]\n",
    "        \n",
    "        prompt_q = '' \n",
    "        prompt_q += 'Question: ' + q + '\\n'\n",
    "        prompt_q += 'Options: ' + c + '\\n'\n",
    "        prompt_q += \"Let's think step by step\\n\"\n",
    "        \n",
    "        prompt_recall = construct_prompt_recall(prompt_paraphrase, prompt_complex, q)\n",
    "        fd_prompts.write(prompt_recall + '\\n----\\n\\n')\n",
    "        \n",
    "        prompt_full = prompt_recall + prompt_q\n",
    "        response = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                            prompt=prompt_full, \n",
    "                                            temperature=0, \n",
    "                                            max_tokens=256)\n",
    "        \n",
    "        ans_model = response['choices'][0]['text']\n",
    "        fd.write('%sA_model:\\n%s\\nA:\\n%s\\n\\n' % (prompt_q, ans_model, a))\n",
    "        # i += 1\n",
    "        # if(i == 2): break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "d7272205-ca6a-45af-a360-e3e55da94c9b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_q 200 correct 47 ratio 0.2350\n"
     ]
    }
   ],
   "source": [
    "_, _, _ =parse_pred_ans('outputs/dev_outputs_recall.txt')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
